# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/8 11:25 
@Author  : ZhangShenao 
@File    : 2.使用ConversationSummaryBufferMemory实现摘要缓冲混合记忆.py 
@Desc    : 使用ConversationSummaryBufferMemory实现摘要缓冲混合记忆
"""
import os
from operator import itemgetter

import dotenv
from langchain.memory import ConversationSummaryBufferMemory
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import MessagesPlaceholder, ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_openai import ChatOpenAI

# 系统指令
SYSTEM_PROMPT = '你是一个具有记忆功能的AI助手，请根据结合用户的当前提问，和之前的历史聊天记录，生成准确的回答。你们的对话全部采用中文'

# 调用的模型名称
MODEL_NAME = 'gpt-3.5-turbo-16k'

# 聊天历史Key
HISTORY_KEY = 'history'

# 加载环境变量
dotenv.load_dotenv()

# 构建LLM
llm = ChatOpenAI(model=MODEL_NAME, openai_api_base=os.getenv('OPENAI_API_BASE'))

# 创建Memory记忆组件
# 使用ConversationSummaryBufferMemory组件,实现摘要缓冲混合记忆
memory = ConversationSummaryBufferMemory(
    input_key='input',
    output_key='output',
    memory_key=HISTORY_KEY,
    max_token_limit=200,  # 限制记忆缓冲的最大token数量
    llm=llm,  # 传入生成摘要是需要使用的LLM
    return_messages=True,  # 结果返回聊天消息列表
)

# 编排Prompt
prompt = ChatPromptTemplate.from_messages([
    ('system', SYSTEM_PROMPT),
    MessagesPlaceholder(HISTORY_KEY),  # 使用MessagesPlaceholder消息占位符
    ('human', '{input}')
])

# 创建OutputParser
parser = StrOutputParser()

# 构建Chain
# 使用RunnablePassthrough.assign()函数,为prompt添加新值
# 使用RunnableLambda将一个lambda表达式封装成Runnable可运行协议
# 使用itemgetter()函数,从memory返回的dict中获取history的值,返回一个字符串
chain = RunnablePassthrough.assign(
    history=RunnableLambda(memory.load_memory_variables) | itemgetter(HISTORY_KEY)
) | prompt | llm | parser

# 获取用户输入
while True:
    input_str = input('Human: ')
    if 'bye' == input_str:
        print('bye bye~')
        break

    # 流式执行Chain
    response = chain.stream({'input': input_str})
    print('AI：', flush=True, end='')
    output = ''
    for trunk in response:
        print(f'{trunk}', flush=True, end='')
        output += trunk

    # 保存当前轮次对话到Memory中
    memory.save_context(inputs={'input': input_str}, outputs={'output': output})

    # 打印当前对话历史
    print(f'\n当前对话历史: {memory.load_memory_variables({})}')
